import numpy as np
from typing import List, Tuple, Optional
from dataclasses import dataclass
from fg_grc_plate import GrapheneReinforcedComposite, MaterialProperties, DistributionPattern

@dataclass
class Layer:
    """Individual layer in the laminated plate."""
    thickness: float
    material: MaterialProperties
    angle: float = 0.0  # Fiber orientation angle in degrees
    is_grc: bool = False
    grc_material: Optional[GrapheneReinforcedComposite] = None
    grc_pattern: DistributionPattern = DistributionPattern.UD
    grc_volume_fraction: float = 0.0
    grc_exponent: float = 1.0

class LaminatedPlate:
    """
    Laminated composite plate with FSDT formulation.
    Based on the theory from the paper.
    """
    
    def __init__(self, layers: List[Layer], length: float, width: float):
        self.layers = layers
        self.length = length  # Plate length (a)
        self.width = width    # Plate width (b)
        self.total_thickness = sum(layer.thickness for layer in layers)
        self.n_layers = len(layers)
        
        # Calculate layer coordinates
        self.layer_coordinates = self._calculate_layer_coordinates()
        
    def _calculate_layer_coordinates(self) -> List[Tuple[float, float]]:
        """Calculate z-coordinates for each layer."""
        coordinates = []
        z = -self.total_thickness / 2
        
        for layer in self.layers:
            z_bottom = z
            z_top = z + layer.thickness
            coordinates.append((z_bottom, z_top))
            z = z_top
            
        return coordinates
    
    def get_material_properties_at_z(self, layer_idx: int, z: float) -> MaterialProperties:
        """Get material properties at specific z-coordinate within a layer."""
        layer = self.layers[layer_idx]
        
        if not layer.is_grc:
            return layer.material
            
        # For GRC layer, calculate volume fraction at z
        z_bottom, z_top = self.layer_coordinates[layer_idx]
        V_G = layer.grc_material.volume_fraction(
            z, z_bottom, z_top, 
            layer.grc_volume_fraction, 
            layer.grc_exponent,
            layer.grc_pattern
        )
        
        return layer.grc_material.effective_properties(V_G)
    
    def get_reduced_stiffness_at_z(self, layer_idx: int, z: float) -> np.ndarray:
        """Get reduced stiffness matrix at specific z-coordinate."""
        layer = self.layers[layer_idx]
        
        if not layer.is_grc:
            # For regular layers, use constant properties
            props = layer.material
            Q11 = props.E / (1 - props.nu**2)
            Q22 = Q11
            Q12 = props.nu * props.E / (1 - props.nu**2)
            Q66 = props.E / (2 * (1 + props.nu))
            Q44 = Q66
            Q55 = Q66
        else:
            # For GRC layers, calculate at specific z
            z_bottom, z_top = self.layer_coordinates[layer_idx]
            V_G = layer.grc_material.volume_fraction(
                z, z_bottom, z_top,
                layer.grc_volume_fraction,
                layer.grc_exponent, 
                layer.grc_pattern
            )
            return layer.grc_material.reduced_stiffness_matrix(V_G)
        
        Q = np.zeros((5, 5))
        Q[0, 0] = Q11
        Q[1, 1] = Q22
        Q[0, 1] = Q[1, 0] = Q12
        Q[2, 2] = Q66
        Q[3, 3] = Q44
        Q[4, 4] = Q55
        
        return Q
    
    def calculate_laminate_stiffness(self, n_integration_points: int = 10) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Calculate laminate stiffness matrices A, B, D.
        
        Args:
            n_integration_points: Number of integration points per layer
            
        Returns:
            A, B, D matrices (3x3 each for in-plane, coupling, bending)
        """
        A = np.zeros((3, 3))
        B = np.zeros((3, 3))
        D = np.zeros((3, 3))
        
        for layer_idx, layer in enumerate(self.layers):
            z_bottom, z_top = self.layer_coordinates[layer_idx]
            
            # Integration points within the layer
            z_points = np.linspace(z_bottom, z_top, n_integration_points)
            dz = (z_top - z_bottom) / (n_integration_points - 1)
            
            for z in z_points:
                Q = self.get_reduced_stiffness_at_z(layer_idx, z)
                
                # Extract in-plane stiffness terms
                Q_reduced = Q[:3, :3]
                
                # Numerical integration
                weight = dz if z != z_points[0] and z != z_points[-1] else dz/2
                
                A += Q_reduced * weight
                B += Q_reduced * z * weight
                D += Q_reduced * z**2 * weight
        
        return A, B, D
    
    def calculate_shear_stiffness(self, n_integration_points: int = 10) -> np.ndarray:
        """
        Calculate transverse shear stiffness matrix.
        
        Returns:
            2x2 shear stiffness matrix
        """
        A_s = np.zeros((2, 2))
        k_s = 5/6  # Shear correction factor
        
        for layer_idx, layer in enumerate(self.layers):
            z_bottom, z_top = self.layer_coordinates[layer_idx]
            
            # Integration points within the layer
            z_points = np.linspace(z_bottom, z_top, n_integration_points)
            dz = (z_top - z_bottom) / (n_integration_points - 1)
            
            for z in z_points:
                Q = self.get_reduced_stiffness_at_z(layer_idx, z)
                
                # Extract shear terms
                Q_shear = Q[3:5, 3:5]
                
                # Numerical integration
                weight = dz if z != z_points[0] and z != z_points[-1] else dz/2
                
                A_s += Q_shear * weight
        
        return k_s * A_s
    
    def calculate_mass_inertia(self, n_integration_points: int = 10) -> Tuple[float, float, float]:
        """
        Calculate mass inertia terms I0, I1, I2.
        
        Returns:
            I0, I1, I2: Mass inertia terms
        """
        I0 = I1 = I2 = 0.0
        
        for layer_idx, layer in enumerate(self.layers):
            z_bottom, z_top = self.layer_coordinates[layer_idx]
            
            # Integration points within the layer
            z_points = np.linspace(z_bottom, z_top, n_integration_points)
            dz = (z_top - z_bottom) / (n_integration_points - 1)
            
            for z in z_points:
                props = self.get_material_properties_at_z(layer_idx, z)
                rho = props.rho
                
                # Numerical integration
                weight = dz if z != z_points[0] and z != z_points[-1] else dz/2
                
                I0 += rho * weight
                I1 += rho * z * weight
                I2 += rho * z**2 * weight
        
        return I0, I1, I2

class FSDTKinematics:
    """
    First-order shear deformation theory kinematics.
    """
    
    @staticmethod
    def strain_displacement_relations(u0, v0, w0, phi_x, phi_y, x, y):
        """
        Calculate strain components from displacement field.
        
        Args:
            u0, v0, w0: Mid-plane displacements
            phi_x, phi_y: Rotations
            x, y: Coordinates
            
        Returns:
            epsilon0, kappa, gamma0: Strain components
        """
        # Mid-plane strains
        epsilon0 = np.array([
            np.gradient(u0, x, axis=1),     # du0/dx
            np.gradient(v0, y, axis=0),     # dv0/dy
            np.gradient(u0, y, axis=0) + np.gradient(v0, x, axis=1)  # du0/dy + dv0/dx
        ])
        
        # Curvatures
        kappa = np.array([
            np.gradient(phi_x, x, axis=1),  # dphi_x/dx
            np.gradient(phi_y, y, axis=0),  # dphi_y/dy
            np.gradient(phi_x, y, axis=0) + np.gradient(phi_y, x, axis=1)  # dphi_x/dy + dphi_y/dx
        ])
        
        # Transverse shear strains
        gamma0 = np.array([
            phi_y + np.gradient(w0, y, axis=0),  # phi_y + dw0/dy
            phi_x + np.gradient(w0, x, axis=1)   # phi_x + dw0/dx
        ])
        
        return epsilon0, kappa, gamma0

# Example usage
if __name__ == "__main__":
    from fg_grc_plate import GrapheneReinforcedComposite, MaterialProperties, GrapheneGeometry
    
    # Material properties
    epoxy = MaterialProperties(E=3.0, nu=0.34, rho=1200)
    graphene = MaterialProperties(E=1010, nu=0.186, rho=1060)
    gpl_geometry = GrapheneGeometry()
    
    # Create GRC material
    grc = GrapheneReinforcedComposite(epoxy, graphene, gpl_geometry)
    
    # Define layers for [0°/GRC/90°] laminate
    layers = [
        Layer(thickness=0.5e-3, material=epoxy, angle=0.0),
        Layer(thickness=1.0e-3, material=epoxy, angle=0.0, is_grc=True, 
              grc_material=grc, grc_pattern=DistributionPattern.UD, 
              grc_volume_fraction=0.01, grc_exponent=1.0),
        Layer(thickness=0.5e-3, material=epoxy, angle=90.0)
    ]
    
    # Create laminated plate
    plate = LaminatedPlate(layers, length=0.1, width=0.1)  # 10cm x 10cm
    
    # Calculate stiffness matrices
    A, B, D = plate.calculate_laminate_stiffness()
    A_s = plate.calculate_shear_stiffness()
    I0, I1, I2 = plate.calculate_mass_inertia()
    
    print("Laminated plate analysis:")
    print(f"Total thickness: {plate.total_thickness*1000:.2f} mm")
    print(f"Number of layers: {plate.n_layers}")
    print(f"A11 = {A[0,0]:.2e} N/m")
    print(f"D11 = {D[0,0]:.2e} N⋅m")
    print(f"I0 = {I0:.2e} kg/m²")
    print(f"Shear stiffness A44 = {A_s[0,0]:.2e} N/m")